package thired.class10_graph;

import java.util.*;

public class Class07_Kruskal {
    public static class UnionFind{
        public HashMap<Node,Node> parent;
        public HashMap<Node,Integer> size;

        public UnionFind() {
            this.parent = new HashMap<>();
            this.size = new HashMap<>();
        }
        public Node find(Node o1){
            Stack<Node> stack = new Stack<>();
            while (o1 != parent.get(o1)){
                stack.add(o1);
                o1 = parent.get(o1);
            }
            while (!stack.isEmpty()){
                parent.put(stack.pop(),o1);
            }
            return o1;
        }


        public void makeSets(Collection<Node> nodes) {
            parent.clear();
            size.clear();
            for (Node node : nodes) {
                parent.put(node, node);
                size.put(node, 1);
            }
        }

        public boolean isSameUnion(Node f1,Node f2){
            return find(f1) == find(f2);
        }

        public void union(Node a, Node b) {
            if (a == null || b == null) {
                return;
            }
            Node aDai = find(a);
            Node bDai = find(b);
            if (aDai != bDai) {
                int aSetSize = size.get(aDai);
                int bSetSize = size.get(bDai);
                if (aSetSize <= bSetSize) {
                    parent.put(aDai, bDai);
                    size.put(bDai, aSetSize + bSetSize);
                    size.remove(aDai);
                } else {
                    parent.put(bDai, aDai);
                    size.put(aDai, aSetSize + bSetSize);
                    size.remove(bDai);
                }
            }
        }
    }


    public static Set<Edge> kruskalMST(Graph graph){
        UnionFind unionFind = new UnionFind();
        unionFind.makeSets((Collection<Node>) graph.nodes);
        HashSet<Edge> res = new HashSet<>();
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>((o1,o2) ->{return  o1.value - o2.value;});
        priorityQueue.addAll(graph.edges);
        while (!priorityQueue.isEmpty()){
            Edge cur = priorityQueue.poll();
            if(!unionFind.isSameUnion(cur.from, cur.to)){
                res.add(cur);
                unionFind.union(cur.from,cur.to);
            }
        }
        return res;
    }
}
